#!/opt/local/bin/python

import csv
import matplotlib.pyplot as plt
from numpy import array

class datareader(object):
    def __init__(self, fname, firstrow=2):
        self.columns = {'ncells': 1,
                'beta0': 2,
                'EoT': 3,
                'beta': 6,
                'Tmid': 7,
                'kTpmid': 8,
                'k2Tppmid': 9,
                'Tin': 10,
                'kTpin': 11,
                'k2Tppin': 12,
                'Tout': 13,
                'kTpout': 14,
                'k2Tppout': 15,
                'phase': 17}
        A = csv.reader(open(fname))

        self.data = {}
        for counter, row in enumerate(A):
            if counter<firstrow-1: continue
            else:
                row = ' '.join(row).split()
                for dat in self.columns:
                    try:
                        self.data[dat].append( float(row[self.columns[dat]]) )
                    except KeyError:
                        self.data[dat] = [float(row[self.columns[dat]])]
        for i in self.data:
            self.data[i] = array(self.data[i])

    def calcT(self):
        kappa = self.data['beta'] / self.data['beta0']

        Tin  = self.data[ 'Tin' ]
        Tmid = self.data[ 'Tmid' ]
        Tout = self.data[ 'Tout' ]

        kTpin  = self.data[ 'kTpin' ]
        kTpmid = self.data[ 'kTpmid' ]
        kTpout = self.data[ 'kTpout' ]

        k2Tppin  = self.data[ 'k2Tppin' ]
        k2Tppmid = self.data[ 'k2Tppmid' ]
        k2Tppout = self.data[ 'k2Tppout' ]

        self.T = fiveCellAv(Tin, Tmid, Tout) + fiveCellAv(kTpin, kTpmid, kTpout)*(1-kappa) + 0.5*fiveCellAv(k2Tppin, k2Tppmid, k2Tppout)*(1-kappa)**2

    def calcE0(self):
        self.E0T = self.data['EoT']
        self.E0  = self.data['EoT'] / self.T

def fiveCellAv(a, b, c):
    return (a + 3*b + c)/5

if __name__=='__main__':
    A = datareader(fname = '/Users/stephenmolloy/Documents/RFVectors/scLinacDetails.csv')
    A.calcT()
    A.calcE0()

    fig1 = plt.figure(1)
    ax = fig1.add_subplot(111)
    ax.plot(A.E0/1e6, 'x')
    ax.set_ylim(bottom=0)
    ax.axvline(x=36-0.5, color='r')
    ax.axvline(x=100-0.5, color='r')
    ax.grid()
    ax.set_xlabel('Cavity #')
    ax.set_ylabel('Peak E-field (MV/m)')
    plt.show()

